package leetcode

import (
	"container/heap"

	"gitee.com/ombak/leetcode/pkg/list"
)

type ItemHeap []*list.ListNode

func (h ItemHeap) Less(i, j int) bool {
	return h[i].Val < h[j].Val
}

func (h ItemHeap) Len() int {
	return len(h)
}

func (h ItemHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *ItemHeap) Push(x interface{}) {
	*h = append(*h, x.(*list.ListNode))
}

func (h *ItemHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// https://leetcode.com/problems/merge-k-sorted-lists/
// power by container/heap.
func mergeKLists(lists []*list.ListNode) *list.ListNode {
	pq := &ItemHeap{}
	heap.Init(pq)
	for _, l := range lists {
		if l != nil {
			heap.Push(pq, l)
		}
	}
	dummy := &list.ListNode{}
	p := dummy
	for pq.Len() > 0 {
		elem := heap.Pop(pq).(*list.ListNode)
		if elem.Next != nil {
			heap.Push(pq, elem.Next)
		}
		elem.Next = nil
		p.Next = elem
		p = p.Next
	}
	return dummy.Next
}

// power by mergeTwoSortedList.
func mergeKLists2(lists []*list.ListNode) *list.ListNode {
	n := len(lists)
	if n < 1 {
		return nil
	}
	for i := 1; i < n; i++ {
		lists[i] = mergeTwoLists(lists[i-1], lists[i])
	}
	return lists[n-1]
}

func mergeTwoLists(l1, l2 *list.ListNode) *list.ListNode {
	dummy := &list.ListNode{}
	p := dummy
	for l1 != nil && l2 != nil {
		if l1.Val <= l2.Val {
			p.Next = l1
			l1 = l1.Next
		} else {
			p.Next = l2
			l2 = l2.Next
		}
		p = p.Next
	}
	if l1 != nil {
		p.Next = l1
	} else {
		p.Next = l2
	}
	return dummy.Next
}
